﻿<!DOCTYPE html>
<html lang="en">

<head>
    <!-- Google tag (gtag.js) -->
    <script async src="https://www.googletagmanager.com/gtag/js?id=G-LY5RQBNBBW"></script>
    <script>
    window.dataLayer = window.dataLayer || [];
    function gtag(){dataLayer.push(arguments);}
    gtag('js', new Date());

    gtag('config', 'G-LY5RQBNBBW');
    </script>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Neural Network Demo</title>

    <script type="module" src="index.js"></script>

    <link rel="stylesheet" href="demo.css">
    <link rel="icon" type="image/x-icon" href="/assets/icon.png">
</head>

<body>
    <script src="bundle.js"></script>
    <script type="module">
        import { torch } from './index.js';
        import { train, test } from './data/data.js';
        window.train = train;
        window.test = test;
        window.torch = torch;
    </script>

    <div style="height: 20px;"></div>
    <div class="container">
        <div class="separator" id="right-separator" style="display: block; padding-bottom: 0px; margin-bottom: 15px;">
            <a target="_blank" href="https://github.com/eduardoleao052/js-pytorch"><img style="display: block;" class='icon' src="../../assets/demo_logo.png"></a>
            <a target="_blank" href="https://github.com/eduardoleao052/js-pytorch"> <img class="logos" src="../../assets/github.png" style="width: 32px; height: 32px; display: inline-block; margin-bottom: 5px;"></a>
        </div>
    
        <div class="separator" id="right-separator" style="margin-top: 10px; display: block;">
            <p style="display: block; font-size: 16px; color: #292929; margin-bottom: 25px;">Welcome to <b>JS-PyTorch's Web Demo</b>! You can set the <b>Model Layers</b> on the right (number of layers and hidden dimension of each layer), and choose the <b>Model Hyperparameters</b> in <u>real time</u>. </p>
            <div id="mnist-and-list">
                <ul> 
                    <li>This model was trained on <b>MNIST</b>, using <b>2000</b> training and <b>100</b> validation images. </li>
                    <br>
                    <li>The images have a dimension of <b>28x28</b>. </li>
                    <br>
                    <li>The output is a number from 0 to 9.</li>
                    <br>
                    <li>The training uses the <b>Adam Optimizer</b> and a <b>Cross Entropy Loss</b>.</li>
                </ul>
                <div id="mnist-div">
                    <img id='mnist' src="data/image1.png">
                </div>
            </div>
        </div>
    </div>

    <div class="container" id="upper-container">
        <div class="left-div">
            <h2 style="display: inline-block; margin-left: 20px; margin-top: 2px; margin-bottom: 10px;">Training Hyperparameters</h2>
            <div class="separator" id="left-separator">

                <label for="batch-size">Batch Size:</label>
                <input type="number" id="batch-size" name="batch-size" min="1" value="32"> 

                <label for="learning-rate">Learning Rate:</label>
                <input type="number" id="learning-rate" name="learning-rate" step="0.00001" min="0" value="0.0001" required>

                <label for="regularization">L2 Regularization:</label>
                <input type="number" id="regularization" name="regularization" step="0.001" min="0" value="0.001">

                <label for="beta1">Beta 1 (Adam):</label>
                <input type="number" id="beta1" name="beta1" step="0.001" min="0" value="0.99">

                <label for="beta2">Beta 2 (Adam):</label>
                <input type="number" id="beta2" name="beta2" step="0.0001" min="0" value="0.999">

            </div>
        </div>
        <div class="right-div">
            <div>
                <h2 style="display: inline-block; margin-right: 15px; margin-left: 20px;">Model Layers</h2>
                <button class='layer-button' onclick="addBox()">+</button>
                <button class='layer-button' onclick="removeBox()">-</button>
                <div class="separator" id="layersBox">
                    
                </div>
            </div>
            <div id="training-buttons">
                <button id="start-button" onclick="trainLoopInitializer()">Train</button>
                <button id="stop-button" onclick="pauseTraining()">Pause</button>
                <button id="reset-button" onclick="resetTraining()">Reset</button>
                <div style="width: 25px; display: inline-block;"></div>
                <button id="cpu-trigger" class="device-button" onclick="CPUCaller()">CPU</button>
                <button id="gpu-trigger" class="device-button" onclick="GPUCaller()">GPU</button>
            </div>
        </div>
    </div>

    <div class="container" style="padding-bottom: 0px;">
        <h2>Graph</h2>
        <!-- This is where the graph will be displayed -->
        <canvas id="graph" width="700" height="300"></canvas>
        <div style="display: inline-block; width: 90%; margin-bottom: 15px;">
            <p id="loss"> <b>Validation Loss:</b> </p>
            <p id="acc"> <b>Validation Accuracy:</b> </p>
            <p id="iter"> <b>Iteration:</b> </p>
            <p id="epoch"> <b>Epoch:</b> </p>
            <p id="total-visited"> <b>Total Training Examples:</b> </p>
            <p id="device-showcase"> <b>Device:</b> </p>
        </div>
        <divs tyle="display: inline-block;">
            <a target="_blank" href="https://github.com/eduardoleao052/js-pytorch"><img class="logos" src="../../assets/github.png"  style="width: 25px; height: 25px; display: inline-block; margin-bottom: 0px;"></a>
            <a target="_blank" href="https://www.linkedin.com/in/eduardo-leao-368bb6259/"><img class="logos" src="../../assets/linkedin.png"  style="width: 25px; height: 25px; display: inline-block; margin-bottom: 0px;"></a>
        </div>
    </div>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/BrowserFS/2.0.0/browserfs.min.js"></script>
    <script src="./gpu-browser.min.js"></script>
    <script>
        let boxCount = [];
        let data = [];
        let training = false;
        let in_loop = true;
        let overFlow = 1;
        let iter = 0;
        let total_visited = 0;
        let device = 'cpu';
        let smoothAcc = 0.1;

        function GPUCaller() {
            if (!training) {
                device = 'gpu';
                let _cpu = document.getElementById('cpu-trigger');
                let _gpu = document.getElementById('gpu-trigger');
                _gpu.style.backgroundColor = '#3e3e3e';
                _cpu.style.backgroundColor = '';
            }
        }
        function CPUCaller() {
            if (!training) {
                device = 'cpu';
                let _cpu = document.getElementById('cpu-trigger');
                let _gpu = document.getElementById('gpu-trigger');
                _cpu.style.backgroundColor = '#3e3e3e';
                _gpu.style.backgroundColor = '';
            }
        }


        function addBox() {
            if (boxCount.length < 5 && !training) {
                // Find container to add Boxes in:
                const container = document.getElementById('layersBox');
                // Create the box:
                const newBox = document.createElement('div');
                newBox.id = `box-div-${boxCount.length}`
                newBox.classList.add('box');
                // Create a numeric input field inside the box (for number of neurons):
                const inputField = document.createElement('input');
                inputField.type = 'number';
                // The next Box has as many neurons as last current Box:
                inputField.value = boxCount[boxCount.length -1] || '64';
                inputField.idx = boxCount.length;
                inputField.id = `box-input`;
                inputField.onchange = function() {
                    changeDim(this, boxCount);
                };
                newBox.appendChild(inputField);
                container.appendChild(newBox);
                boxCount.push(Number(inputField.value));
            };
        };

        function removeBox() {
            if (boxCount.length > 1 && !training) {
                const container = document.getElementById('layersBox');
                container.removeChild(container.lastChild);
                boxCount.pop();
            };
        };

        function changeDim(el, boxCount) {
            boxCount[el.idx] = Number(el.value);                
        };

        function getBatch(data, batch_size) {
            // Instantiate x_batch and y_batch as empty tensors:
            let x_batch = [];
            let y_batch = [];
            // Iteratively add instances to batch:
            for (let i=0 ; i < batch_size ; i++) {
                p = Math.floor(Math.random() * data.length);
                x_batch.push(Object.entries(data[p])[0][1])
                y_batch.push(Number(Object.entries(data[p])[0][0]))
            };
            return [torch.tensor(x_batch), torch.tensor(y_batch)];
        };

        function changeMNIST() {
            let p = Math.floor(Math.random()*20) + 1;
            document.getElementById("mnist").src=`data/image${p}.png`;
        };

        function trainLoopInitializer() {
            in_loop = true;
            trainLoop();
        };

        function trainLoop() {
            if (!in_loop) return;
            trainStep();
            setTimeout(trainLoop, 0.01);
        };
       
        function trainStep() {
            if (!training) {
                training = true;
                buttons = document.getElementsByClassName('layer-button');
                for (button of buttons) {
                    button.style.backgroundColor = '#0056b3';
                };
                let _cpu = document.getElementById('cpu-trigger');
                let _gpu = document.getElementById('gpu-trigger');
                if (device === 'cpu') {
                    _cpu.style.backgroundColor = '#3e3e3e';
                } else {
                    _gpu.style.backgroundColor = '#3e3e3e';
                }
                // Implement dummy torch.nn.Module class:
                class NeuralNet extends torch.nn.Module {
                    constructor() {
                        super();
                        // Instantiate Neural Network's Layers:
                        this.wIn = new torch.nn.Linear(784,boxCount[0],device);
                        this.reluIn = new torch.nn.ReLU();
                        if (boxCount.length > 1) {
                            this.w1 = new torch.nn.Linear(boxCount[0],boxCount[1],device);
                            this.relu1 = new torch.nn.ReLU();
                        } if (boxCount.length > 2) {
                            this.w2 = new torch.nn.Linear(boxCount[1],boxCount[2],device);
                            this.relu2 = new torch.nn.ReLU();
                        } if (boxCount.length > 3) {
                            this.w3 = new torch.nn.Linear(boxCount[2],boxCount[3],device);
                            this.relu3 = new torch.nn.ReLU();
                        } if (boxCount.length > 4) {
                            this.w4 = new torch.nn.Linear(boxCount[3],boxCount[4],device);
                            this.relu4 = new torch.nn.ReLU();
                        };
                        this.wOut = new torch.nn.Linear(boxCount[boxCount.length-1], 10, device);
                    
                    };

                    forward(x) {
                        let z;
                        z = this.wIn.forward(x);
                        z = this.reluIn.forward(z);
                        
                        if (boxCount.length > 1) {
                            z = this.w1.forward(z);
                            z = this.relu1.forward(z);
                        } if (boxCount.length > 2) {
                            z = this.w2.forward(z);
                            z = this.relu2.forward(z);
                        } if (boxCount.length > 3) {
                            z = this.w3.forward(z);
                            z = this.relu3.forward(z);
                        } if (boxCount.length > 4) {
                            z = this.w4.forward(z);
                            z = this.relu4.forward(z);
                        };
                        z = this.wOut.forward(z);
                        z = z.div(100);

                        return z;
                    };
                };
                globalThis.model = new NeuralNet();

                // Define loss function and optimizer:
                globalThis.loss_func = new torch.nn.CrossEntropyLoss();
            }; 

            // Get live learning rate and regularization values.
            let batch_size = Number(document.getElementById('batch-size').value)
            let lr = Number(document.getElementById('learning-rate').value)
            let reg = Number(document.getElementById('regularization').value )
            let beta1 = Number(document.getElementById('beta1').value )
            let beta2 = Number(document.getElementById('beta2').value )
            let eps = Number(0.000001)

            // Build optimizer:
            let optimizer = new torch.optim.Adam(model.parameters(), lr=lr, reg=reg, betas=[beta1, beta2], eps=eps)
            let loss;
            let loss_test;

            // Training Loop:
            for(let i=0 ; i < 1 ; i++) {
                let [x_batch, y_batch] = getBatch(train, batch_size)
                
                let z = model.forward(x_batch)

                // Get loss:
                loss = loss_func.forward(z, y_batch)

                // Backpropagate the loss using neuralforge.tensor's backward() method:
                loss.backward()

                // Update the weights:
                optimizer.step()
                
                // Reset the gradients to zero after each training step:
                optimizer.zero_grad()
                
                // Now, get a batch of test data:
                let [x_test, y_test] = getBatch(test, batch_size)

                // Run the batch of test data through the model:
                let z_test = model.forward(x_test)

                // Get the test loss and accuracy:
                acc_test = getAccuracy(z_test, y_test)
                loss_test = loss_func.forward(z_test, y_test)
                smoothAcc = acc_test * 0.02 + smoothAcc * 0.98

                // If loss went to infinity (model way too large for training size), represent that in the graph:
                if (isNaN(loss_test.data[0])) {
                    overFlow = overFlow * 1.5;
                    data.push(overFlow + (Math.random() - 0.5) * 15);
                // If not, just keep adding the loss to the graph:
                } else {data.push(loss_test.data)}

                // Display iteration and loss on the screen:
                document.getElementById('iter').innerHTML = `<b>Iteration:</b> ${iter}`;
                document.getElementById('total-visited').innerHTML = `<b>Total Training Examples:</b> ${total_visited}`;
                document.getElementById('loss').innerHTML = `<b>Validation Loss:</b> ${loss_test.data[0].toFixed(3)}`;
                document.getElementById('device-showcase').innerHTML = `<b>Device:</b> ${device}`
                document.getElementById('epoch').innerHTML = `<b>Epoch:</b> ${Math.floor(total_visited/train.length)}`;
                document.getElementById('acc').innerHTML = `<b>Validation Accuracy:</b> ${smoothAcc.toFixed(3)}`;
                iter += 1;
                total_visited += batch_size;
            
                plotGraph()

                optimizer.zero_grad()
            };
        };

        function getAccuracy(z, y) {
            const [B, D] = z.shape;
            let acc = 0;
            for (let i = 0; i < B; i+=1) {
                let biggest = 0;
                let biggestIndex = 0;
                for (let j = 0; j < D; j+=1) {
                    if (z.data[i][j] > biggest) {
                        biggestIndex = j;
                        biggest = z.data[i][j];
                    }
                }
                if (biggestIndex === y.data[i]) {acc += 1}
            }
            return acc / B
        }

        function pauseTraining() {
            in_loop = false;
        };

        function resetTraining() {
            in_loop = false;
            training = false;
            smoothAcc = 0.1;
            iter = 0;
            total_visited = 0;
            data = [];
            plotGraph();
            buttons = document.getElementsByClassName('layer-button');
            for (button of buttons) {
                button.style.backgroundColor = '';

            }
            buttons = document.getElementsByClassName('device-button');
            for (button of buttons) {
                button.style.backgroundColor = '';
            }
        };

        function plotGraph() {
            var canvas = document.getElementById('graph');
            var ctx = canvas.getContext('2d');
            ctx.clearRect(0, 0, canvas.width, canvas.height);

            var startX = 50;
            var startY = canvas.height - 30;
            var maxX = data.length - 1;
            var maxY = Math.max(...data);
            var stepX = (canvas.width - 2 * startX) / maxX;
            var stepY = (startY - 30) / maxY;

            // Set default stroke style and line width
            ctx.strokeStyle = 'black';
            ctx.lineWidth = 1;

            // Draw x-axis label
            ctx.fillStyle = 'black';
            ctx.font = '12px Arial';
            ctx.textAlign = 'center';
            ctx.fillText('Iterations', canvas.width / 2, startY +15);

            // Draw y-axis label
            ctx.save();
            ctx.translate(8, canvas.height / 2);
            ctx.rotate(-Math.PI / 2);
            ctx.fillStyle = 'black';
            ctx.font = '12px Arial';
            ctx.textAlign = 'center';
            ctx.fillText('Cross-Entropy Loss', 0, 5);
            ctx.restore();

            // Draw grid lines
            ctx.beginPath();
            for (var i = 1; i < 5; i += 1) {
                var y = (startY * i )/ 5
                ctx.moveTo(startX, y);
                ctx.lineTo(canvas.width - startX, y);
            };
            ctx.stroke();

            // Draw x-axis
            ctx.beginPath();
            ctx.moveTo(startX, startY);
            ctx.lineTo(canvas.width - startX, startY);
            ctx.stroke();

            // Draw y-axis
            ctx.beginPath();
            ctx.moveTo(startX, startY);
            ctx.lineTo(startX, 30);
            ctx.stroke();

            // Draw ticks and labels on x-axis
            for (var i = 0; i <= maxX; i++) {
                var x = startX + i * stepX * 100;
                ctx.beginPath();
                ctx.moveTo(x, startY);
                ctx.lineTo(x, startY + 5);
                ctx.stroke();
                ctx.fillText((i*100).toString(), x - 5, startY + 20);
            };

            // Draw ticks and labels on y-axis
            for (var i = 0; i < 4; i += 1) {
                var y = (startY * i) / 5 + startY/5;
                ctx.beginPath();
                ctx.moveTo(startX, y);
                ctx.lineTo(startX - 5, y );
                ctx.stroke();
                if (stepY != 0) {
                    ctx.fillText((y * maxY / stepY * 2 / 5).toFixed(2), startX - 20, startY - (y) + 5 );
                }
                
            };

            // Draw line plot
            ctx.beginPath();
            ctx.strokeStyle = 'blue';
            ctx.lineWidth = 2;
            ctx.moveTo(startX, startY - data[0] * stepY);
            for (var i = 1; i <= maxX; i++) {
                var x = startX + i * stepX;
                var y = startY - data[i] * stepY;
                ctx.lineTo(x, y);
            };
            ctx.stroke();
        };

        // Initial setup
        addBox(); // Add one box initially
        setInterval(changeMNIST , 1200);
    </script>
</body>

</html>